package org.easy.milvus;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.param.ConnectParam;
import io.milvus.param.R;
import io.milvus.param.dml.InsertParam;
import org.easy.core.util.Lists;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * https://milvus.io/docs/zh/install-java.md
 */
@Configuration
@EnableConfigurationProperties({MilvusProperties.class})
public class MilvusAutoConfiguration {

    @Resource
    private MilvusProperties milvusProperties;

    public static void main(String[] args) {
        EmbeddingClient embeddingClient = new EmbeddingClient("http://127.0.0.1:8001/get_text_embedding");
        MilvusServiceClient milvusServiceClient = new MilvusServiceClient(ConnectParam.newBuilder().withHost("127.0.0.1").withPort(19530).build());

        String c = "public_doc";
        MilvusTemplate template = new MilvusTemplate(milvusServiceClient, embeddingClient);
        template.createDefaultCollection(c);
        ArrayList<String> strings = Lists.newArrayList("美丽的天使", "走在路上");

        List<List<Float>> contentVector = embeddingClient.doEmbedding(strings);

        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("content", strings));
        fields.add(new InsertParam.Field("partition_name", Lists.newArrayList("test", "test")));
        fields.add(new InsertParam.Field("content_vector", contentVector));
        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(template.wrapCollectionName(c))
                .withFields(fields)
                .build();
        R<MutationResult> insert = milvusServiceClient.insert(insertParam);

        List<BaseMilvusDto> test = template.search(c, "test2", "天使", 10);

        System.out.println("success");


    }

    @Bean
    public MilvusServiceClient milvusServiceClient() {
        return new MilvusServiceClient(ConnectParam.newBuilder().withHost(milvusProperties.getIpAddr()).withPort(milvusProperties.getPort()).build());
    }

    @Bean
    public EmbeddingClient embeddingClient() {
        return new EmbeddingClient(milvusProperties.getEmbeddingClientUrl());
    }

    @Bean
    public FileChunkClient fileChunkClient() {
        return new FileChunkClient(milvusProperties.getFileChunkClientUrl());
    }

    @Bean
    public MilvusTemplate milvusTemplate() {
        EmbeddingClient embeddingClient = new EmbeddingClient(milvusProperties.getEmbeddingClientUrl());
        MilvusServiceClient milvusServiceClient = new MilvusServiceClient(ConnectParam.newBuilder().withHost(milvusProperties.getIpAddr()).withPort(milvusProperties.getPort()).build());
        return new MilvusTemplate(milvusServiceClient, embeddingClient);
    }
}
